"""
NaiveBayes
灰度：0 / 1
分类：0-9
入口在最后
"""

# train:60000
# test:10000

# acc: 0.8427
# time: 148s

import pandas as pd
import numpy as np
import time
from collections import Counter


def loadData(fileName):
    data = pd.read_csv(fileName, header=None)
    data = data.values  # we use numpy

    y_label = data[:, 0]
    x_label = data[:, 1:]

    # 为简化计算，将灰度直接二分为0和1
    x_label[x_label < 128] = 0
    x_label[x_label >= 128] = 1

    return x_label, y_label

# 计算先验概率，条件概率


def calcP(x_train, y_train):
    y_class = 10
    x_feature = 784  # 图片为28*28

    # 初始化 Py 10, Pxy (Iy,Ixy为初始计数，P为log计算后的，不改变相对性大小)
    iy = np.array([0]*y_class)
    Py = np.array([0]*y_class)
    ixy = np.zeros((y_class, x_feature, 2))  # 改成 2*782*10更好
    Pxy = np.zeros((y_class, x_feature, 2))  # 10*784*2

    for i in range(len(y_train)):
        iy[y_train[i]] += 1

    # log处理
    Py = np.log(iy/len(y_train))

    # Count(x_i(0|1), y_j) 10*784*2
    for i in range(len(x_train)):
        for j in range(len(x_train[i])):
            ixy[y_train[i]][j][x_train[i][j]] += 1

    # log处理
    # laplace smoothing
    # P(x_j = 1 | y_i) 和 P(x_j = 0 | y_i) 若不简化为 10*784*256
    for i in range(y_class):
        print(i)
        for j in range(x_feature):
            Pxy[i][j][0] = np.log((ixy[i][j][0]+1)/(iy[i]+2))
            Pxy[i][j][1] = np.log((ixy[i][j][1]+1)/(iy[i]+2))

    return Py, Pxy


def NB(Py, Pxy, x):  # x 为单一测试数据
    y_class = 10
    x_feature = 784
    P = [0 for _ in range(y_class)]
    print('P size:', len(P))

    # x_test 1*784
    for i in range(y_class):
        for j in range(x_feature):
            P[i] += Pxy[i][j][x[j]]
        P[i] = P[i] + Py[i]  # 由于log 不使用乘法
    return P.index(max(P))


def test(x_train, y_train, x_test, y_test):
    Py, Pxy = calcP(x_train, y_train)
    acc_num = 0
    acc = 0

    for i in range(len(x_test)):
        y_pred = NB(Py, Pxy, x_test[i])
        if y_pred == y_test[i]:
            acc_num += 1
        print(f'find {i}th data cluster:y_pred={y_pred},y={y_test[i]}')
        print('now_acc=', acc_num / (i + 1))


if __name__ == '__main__':
    start = time.time()

    x_train, y_train = loadData('Mnist/mnist_train.csv')
    x_test, y_test = loadData('Mnist/mnist_test.csv')

    test(x_train, y_train, x_test, y_test)

    end = time.time()

    print('run time:', end-start)
